# main.py
import torch
import copy
import numpy as np
from fl_client import FLClient
from fl_server import FLServer, SecurityException
from privacy_compliance_checker import PrivacyComplianceChecker
from model import SimpleCNN
from dataset import DummyDataset

if __name__ == "__main__":
    # Initialize shared components
    shared_seed = b'fed_learning_secret_key'
    global_model = SimpleCNN()
    server = FLServer(global_model)
    
    # Create simulated clients with data partitions
    train_data = [DummyDataset() for _ in range(5)]
    clients = [
        FLClient(i, train_data[i], copy.deepcopy(global_model), shared_seed)
        for i in range(5)
    ]
    
    # Federated training loop
    total_rounds = 10
    for round_num in range(total_rounds):
        client_updates = []
        quantum_tags = []
        
        # Client local training
        for client in clients:
            noisy_params, q_tag = client.local_train(round_num, total_rounds)
            client_updates.append((noisy_params, q_tag))
            quantum_tags.append(q_tag)
        
        try:
            # Server aggregation
            aggregated_params = server.aggregate_models(client_updates, quantum_tags)
            
            # Noise revocation (using first client's noise generator)
            noise_vector = clients[0].noise_gen.generate_noise(round_num)
            clean_params = server.revoke_global_noise(aggregated_params, noise_vector)
            
            # Update global model
            with torch.no_grad():
                for param, new_val in zip(global_model.parameters(), clean_params):
                    param.copy_(new_val)
            
            # Privacy compliance check
            current_sigma = clients[0].noise_gen.sigma
            epsilon = PrivacyComplianceChecker.calculate_epsilon(current_sigma)
            print(f"Round {round_num}: ε={epsilon:.4f}")
            
            # Dynamic noise adjustment example
            if PrivacyComplianceChecker.check_violation(epsilon):
                print("Privacy violation detected! Adjusting noise parameters")
                for client in clients:
                    client.noise_gen.sigma *= 1.2
        
        except SecurityException as e:
            print(f"Security alert in round {round_num}: {str(e)}")
            # Implement recovery mechanism (e.g., rollback)
            continue
